﻿using StableDiffusionTools.Utils;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net;
using System.Net.Http;
using System.Net.Http.Headers;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;
using System.Text.Json.Serialization;
using System.Threading;
using System.Threading.Tasks;

namespace StableDiffusionTools.Oobabooga
{
    public static class WebUIAPI
    {
        public static string WebUIServer = "127.0.0.1";
        public static int WebUIPort { get; set; }
        public static string OAuthToken { get; set; } = "";
        private static HttpClient streamingClient;
        private static JsonSerializerOptions options = new JsonSerializerOptions { Converters = { new JsonStringEnumConverter(new JsonNamingPolicyLowercase()) } };

        static WebUIAPI()
        {
            streamingClient = new HttpClient();
        }

        public static async Task<RequestResult<Models>> GetModelsInternal(int timeout = 5000)
        {
            return await HttpClientUtil.GetJsonAuth<Models>($"{FormatAPIUrl("internal/model/list")}", OAuthToken, timeout);
        }

        public static async Task<RequestResult<ModelInfo>> GetModelInfoInternal(int timeout = 5000)
        {
            return await HttpClientUtil.GetJsonAuth<ModelInfo>($"{FormatAPIUrl("internal/model/info")}", OAuthToken, timeout);
        }

        public static async Task<RequestResult<string>> LoadModelInternal(string model, int timeout = 5 * 60 * 1000)
        {
            var requestJSON = new JsonObject
            {
                ["model_name"] = model,
                ["args"] = null,
                ["settings"] = null
            };

            var result = await HttpClientUtil.PostJSONAuth($"{FormatAPIUrl("internal/model/load")}", OAuthToken, requestJSON.ToJsonString(), timeout).ConfigureAwait(false);
            string error = null;

            if (result.Data != null)
            {
                bool gotOk = false;
                try
                {
                    string ok = result.Data.Deserialize<string>();
                    gotOk = ok.Equals("OK", StringComparison.InvariantCultureIgnoreCase);
                }
                catch (Exception) { }
                if (!gotOk)
                {
                    try
                    {
                        ErrorResult err = result.Data.Deserialize<ErrorResult>();
                        error = err.detail;
                    }
                    catch (Exception) { }
                }
            }
            result.Data?.Dispose();
            return new RequestResult<string>("", error, error == null);
        }

        public static async Task<RequestResult<string>> UnloadModelInternal(int timeout = 5000)
        {
            var result = await HttpClientUtil.PostJSONAuth($"{FormatAPIUrl("internal/model/unload")}", OAuthToken, null, timeout).ConfigureAwait(false);
            result.Data?.Dispose();
            return new RequestResult<string>("", result.Error, result.Success);
        }

        public static async Task<RequestResult<string>> StopGenerationInternal(int timeout = 5000)
        {
            var result = await HttpClientUtil.PostJSONAuth($"{FormatAPIUrl("internal/stop-generation")}", OAuthToken, null, timeout).ConfigureAwait(false);
            result.Data?.Dispose();
            return new RequestResult<string>("", result.Error, result.Success);
        }

        public static async Task<RequestResult<ChatCompletionResult>> ChatCompletions(List<Message> messages, string model = "", float? temperature = 0.8f, int? maxtokens = -1, float? presence_penalty = 0.0f, float? frequency_penalty = 1.1f, float? top_p = 0.95f, int? top_k = null, int? n = null, int? maxtokenssecond = null, CancellationToken cancel = default(CancellationToken))
        {
            for (int i = 0; i < messages.Count; i++)
            {
                messages[i].FixToApi();
            }
            var chatCompletionRequest = new ChatCompletion()
            {
                messages = messages,
                model = model,
                temperature = temperature,
                max_tokens = maxtokens,
                presence_penalty = presence_penalty,
                frequency_penalty = frequency_penalty,
                top_p = top_p,
                top_k = top_k,
                n = n,
                max_tokens_second = maxtokenssecond
            };

            var result = await HttpClientUtil.PostJSONAuth($"{FormatAPIUrl("chat/completions")}", OAuthToken, JsonSerializer.Serialize(chatCompletionRequest, options), 600000, cancel).ConfigureAwait(false);
            ChatCompletionResult chatCompletion = null;
            if (result.Success && string.IsNullOrEmpty(result.Error))
            {
                chatCompletion = result.Data.Deserialize<ChatCompletionResult>(options);
                for (int i = 0; i < chatCompletion.Choices.Count; i++)
                {
                    chatCompletion.Choices[i].Message.FixFromApi();
                }
            }
            result.Data?.Dispose();

            return new RequestResult<ChatCompletionResult>(chatCompletion, result.Error, result.Success);
        }

        public static async Task<RequestResult<ChatCompletionResult>> ChatCompletionsStreaming(List<Message> messages, string model = "", float? temperature = 0.8f, int? maxtokens = -1, float? presence_penalty = 0.0f, float? frequency_penalty = 1.1f, float? top_p = 0.95f, int? top_k = null, int? n = null, int? maxtokenssecond = null, CancellationToken cancel = default(CancellationToken))
        {
            for (int i = 0; i < messages.Count; i++)
            {
                messages[i].FixToApi();
            }
            var chatCompletionRequest = new ChatCompletion()
            {
                messages = messages,
                model = model,
                temperature = temperature,
                max_tokens = maxtokens,
                presence_penalty = presence_penalty,
                frequency_penalty = frequency_penalty,
                top_p = top_p,
                top_k = top_k,
                n = n,
                max_tokens_second = maxtokenssecond,
                stream = true
            };


            var taskResult = await Task.Run(async () => 
            {
                string error = null;
                RequestResult<ChatCompletionResult> result;

                try
                {
                    using (var requestMessage = new HttpRequestMessage(HttpMethod.Post, FormatAPIUrl("chat/completions")))
                    {
                        requestMessage.Headers.Authorization = new AuthenticationHeaderValue("Bearer", OAuthToken);
                        requestMessage.Content = new StringContent(JsonSerializer.Serialize(chatCompletionRequest, options), Encoding.UTF8, "application/json");
                        using (var cts = new CancellationTokenSource(TimeSpan.FromMilliseconds(600000)))
                        using (CancellationTokenSource linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cts.Token, cancel))
                        using (var response = await streamingClient.SendAsync(requestMessage, HttpCompletionOption.ResponseHeadersRead, linkedCts.Token))
                        {
                            var success = response.StatusCode == HttpStatusCode.OK;
                            var validationError = response.StatusCode == (HttpStatusCode)422;
                            if (!success && !validationError)
                            {
                                result = new RequestResult<ChatCompletionResult>(null, $"Unknown Error: {response.StatusCode}", false);
                            }
                            else
                            {
                                using (var contentStream = await response.Content.ReadAsStreamAsync())
                                using (var streamReader = new StreamReader(contentStream))
                                {
                                    ChatCompletionResult lastChunk = null;
                                    ChatCompletionResult combinedChunks = null;
                                    while (!streamReader.EndOfStream && !linkedCts.Token.IsCancellationRequested)
                                    {
                                        var message = await streamReader.ReadLineAsync();
                                        if(message.StartsWith("data: "))
                                        {
                                            try
                                            {
                                                string json = message.Substring(6);
                                                lastChunk = JsonSerializer.Deserialize<ChatCompletionResult>(json, options);
                                                if (combinedChunks == null)
                                                {
                                                    combinedChunks = new ChatCompletionResult();
                                                    combinedChunks.Choices = new List<Choice>();
                                                    combinedChunks.Choices.Add(lastChunk.Choices[0]);
                                                }
                                                else
                                                {
                                                    var lastChoice = combinedChunks.Choices[combinedChunks.Choices.Count - 1];
                                                    if (lastChoice.Delta.Role == lastChunk.Choices[0].Delta.Role)
                                                    {
                                                        lastChoice.Delta.Content += lastChunk.Choices[0].Delta.Content;
                                                    }
                                                    else
                                                    {
                                                        combinedChunks.Choices.Add(lastChunk.Choices[0]);
                                                    }
                                                }
                                                combinedChunks.Created = lastChunk.Created;
                                                combinedChunks.ID = lastChunk.ID;
                                                combinedChunks.Model = lastChunk.Model;
                                                combinedChunks.@object = lastChunk.@object;
                                                combinedChunks.Usage = lastChunk.Usage;
                                            }
                                            catch (Exception)
                                            {
                                                lastChunk = null;
                                            }
                                        }
                                    }
                                    result = new RequestResult<ChatCompletionResult>(combinedChunks, validationError ? "Validation Error" : "", combinedChunks != null);
                                }
                            }
                        }
                    }
                }
                catch (Exception ex)
                {
                    error = ex.Message;
                    result = new RequestResult<ChatCompletionResult>(null, error, false);
                }

                if(result == null)
                    result = new RequestResult<ChatCompletionResult>(null, "Result was null", false);

                return result;
            });

            ChatCompletionResult chatCompletion = null;
            if (taskResult.Success && string.IsNullOrEmpty(taskResult.Error))
            {
                chatCompletion = taskResult.Data;
                for (int i = 0; i < chatCompletion.Choices.Count; i++)
                {
                    chatCompletion.Choices[i].Delta.FixFromApi();
                }
            }
            //taskResult.Data?.Dispose();

            return new RequestResult<ChatCompletionResult>(chatCompletion, taskResult.Error, taskResult.Success);
        }

        private static string FormatAPIUrl(string apiName)
        {
            return $"http://{WebUIServer}:{WebUIPort}/v1/{apiName}";
        }
    }
}
